#!/usr/bin/env python
# -*- coding: UTF-8 -*-
"""
@Project ：python_learning 
@File ：视频图片截取.py
@IDE  ：PyCharm 
@Author ：李涵彬
@Date ：2025/1/6 22:29 
"""

import cv2
from PIL import Image
import os
from typing import List


class VideoProcessor:
	def __init__(self, video_path: str, output_dir: str, threshold_percent: float, skip_frames: int = 1):
		"""
		初始化视频处理器。

		:param video_path: 视频文件的路径。
		:param output_dir: 输出图片的目录。
		:param threshold_percent: 帧之间差异性的百分比阈值。
		:param skip_frames: 跳帧数，默认为1（不跳帧）。
		"""
		self.video_path = video_path
		self.output_dir = output_dir
		self.threshold_percent = threshold_percent
		self.cap = cv2.VideoCapture(video_path)
		self.prev_frame = None
		self.skip_frames = skip_frames

	def process_video(self):
		"""
		处理视频，识别帧之间的差异性，并存储差异性较大的帧。
		"""
		frame_count = 0
		skip_count = 0
		fps = self.cap.get(cv2.CAP_PROP_FPS)

		# 如果帧率超过15，则设置跳帧数
		if fps > 15:
			self.skip_frames = int(fps // 15)

		while self.cap.isOpened():
			ret, frame = self.cap.read()
			if not ret:
				break

			# 跳帧处理
			if skip_count < self.skip_frames:
				skip_count += 1
				continue
			skip_count = 0

			if self.prev_frame is not None:
				# 计算当前帧和前一帧之间的差异性
				diff = cv2.absdiff(self.prev_frame, frame)
				diff = cv2.cvtColor(diff, cv2.COLOR_BGR2GRAY)
				diff = cv2.threshold(diff, 25, 255, cv2.THRESH_BINARY)[1]
				diff_count = cv2.countNonZero(diff)
				total_pixels = diff.shape[0] * diff.shape[1]
				diff_percent = (diff_count / total_pixels) * 100

				# 如果差异性超过阈值，则存储当前帧
				if diff_percent > self.threshold_percent:
					self.save_frame(frame, frame_count)

			self.prev_frame = frame
			frame_count += 1

		self.cap.release()

	def save_frame(self, frame, frame_count):
		"""
		存储帧为图片。

		:param frame: 要存储的帧。
		:param frame_count: 帧的编号。
		"""
		if not os.path.exists(self.output_dir):
			os.makedirs(self.output_dir)

		img = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
		img.save(os.path.join(self.output_dir, f"frame_{frame_count}.jpg"))


# 示例用法
if __name__ == "__main__":
	video_path = "video.mp4"
	output_dir = "./images"
	threshold_percent = 60.0  # 差异性百分比阈值

	processor = VideoProcessor(video_path, output_dir, threshold_percent)
	processor.process_video()
